package com.ksd.common.core.crypto;

import com.ksd.common.utils.sign.BytesUtils;
import org.bouncycastle.asn1.*;
import org.bouncycastle.crypto.CryptoException;
import org.bouncycastle.crypto.CryptoServicesRegistrar;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.signers.DSAKCalculator;
import org.bouncycastle.crypto.signers.RandomDSAKCalculator;
import org.bouncycastle.math.ec.*;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;

import java.io.IOException;
import java.math.BigInteger;

/**
 * The SM2 Digital Signature algorithm.
 */
public class SM2Signer implements ECConstants {
	private final DSAKCalculator kCalculator = new RandomDSAKCalculator();
	private final SM3Digest digest = new SM3Digest();

	private ECDomainParameters ecParams;
	private ECPoint pubPoint;
	private BigInteger d;
	private byte[] z;

	public void init(boolean forSigning, byte[] ecKey) {
		byte[] userID = Hex.decode("31323334353637383132333435363738"); // the default value

		BigInteger SM2_ECC_P = new BigInteger(Constants.p, 16);
		BigInteger SM2_ECC_A = new BigInteger(Constants.a, 16);
		BigInteger SM2_ECC_B = new BigInteger(Constants.b, 16);
		BigInteger SM2_ECC_N = new BigInteger(Constants.n, 16);
		BigInteger SM2_ECC_GX = new BigInteger(Constants.Gx, 16);
		BigInteger SM2_ECC_GY = new BigInteger(Constants.Gy, 16);

		@SuppressWarnings("deprecation")
		ECCurve curve = new ECCurve.Fp(SM2_ECC_P, SM2_ECC_A, SM2_ECC_B);

		ECPoint g = curve.createPoint(SM2_ECC_GX, SM2_ECC_GY);

		ECDomainParameters domainParams = new ECDomainParameters(curve, g, SM2_ECC_N);
		ecParams = domainParams;

		if (forSigning) {
			d = calculateE(ecKey);

			kCalculator.init(ecParams.getN(), CryptoServicesRegistrar.getSecureRandom());

			pubPoint = createBasePointMultiplier().multiply(ecParams.getG(), d).normalize();
		} else {
			pubPoint = ecParams.getCurve().decodePoint(ecKey);
		}

		z = getZ(userID);
		digest.update(z, 0, z.length);
	}

	public void update(byte b) {
		digest.update(b);
	}

	public void update(byte[] in, int off, int len) {
		digest.update(in, off, len);
	}

	public boolean verifySignature(byte[] signature) {
		try {
//            BigInteger[] rs = derDecode(signature);
//            if (rs != null)
//            {

			byte[] r = new byte[signature.length / 2];

			System.arraycopy(signature, 0, r, 0, r.length);

			byte[] s = new byte[signature.length / 2];

			System.arraycopy(signature, r.length, s, 0, s.length);

			return verifySignature(calculateE(r), calculateE(s));
//            }
		} catch (Exception e) {
		}

		return false;
	}

	public void reset() {
		digest.reset();

		if (z != null) {
			digest.update(z, 0, z.length);
		}
	}

	public byte[] generateSignature() throws CryptoException {
		byte[] eHash = digestDoFinal();

		BigInteger n = ecParams.getN();
		BigInteger e = calculateE(eHash);
		BigInteger d = this.d;

		BigInteger r, s;

		ECMultiplier basePointMultiplier = createBasePointMultiplier();

		// 5.2.1 Draft RFC: SM2 Public Key Algorithms
		do // generate s
		{
			BigInteger k;
			do // generate r
			{
				// A3
				k = kCalculator.nextK();

				// A4
				ECPoint p = basePointMultiplier.multiply(ecParams.getG(), k).normalize();

				// A5
				r = e.add(p.getAffineXCoord().toBigInteger()).mod(n);
			} while (r.equals(ZERO) || r.add(k).equals(n));

			// A6
			BigInteger dPlus1ModN = d.add(ONE).modInverse(n);

			s = k.subtract(r.multiply(d)).mod(n);
			s = dPlus1ModN.multiply(s).mod(n);
		} while (s.equals(ZERO));

		// A7
		try {
//            return derEncode(r, s);

			byte[] br = BytesUtils.asUnsigned32ByteArray(r);
			byte[] bs = BytesUtils.asUnsigned32ByteArray(s);

			byte[] dst = new byte[br.length + bs.length];

			System.arraycopy(br, 0, dst, 0, br.length);
			System.arraycopy(bs, 0, dst, br.length, bs.length);

			return dst;
		} catch (Exception ex) {
			throw new CryptoException("unable to encode signature: " + ex.getMessage(), ex);
		}
	}

	private boolean verifySignature(BigInteger r, BigInteger s) {
		BigInteger n = ecParams.getN();

		// 5.3.1 Draft RFC: SM2 Public Key Algorithms
		// B1
		if (r.compareTo(ONE) < 0) {
			System.out.println(false);
		}
		if (r.compareTo(n) >= 0) {
			System.err.println(false);
		}
		if (r.compareTo(ONE) < 0 || r.compareTo(n) >= 0) {
			return false;
		}

		// B2
		if (s.compareTo(ONE) < 0 || s.compareTo(n) >= 0) {
			return false;
		}

		// B3
		byte[] eHash = digestDoFinal();

		// B4
		BigInteger e = calculateE(eHash);

		// B5
		BigInteger t = r.add(s).mod(n);
		if (t.equals(ZERO)) {
			return false;
		}
		// B6
		ECPoint q = pubPoint;
		ECPoint x1y1 = ECAlgorithms.sumOfTwoMultiplies(ecParams.getG(), s, q, t).normalize();
		if (x1y1.isInfinity()) {
			return false;
		}

		// B7
		BigInteger expectedR = e.add(x1y1.getAffineXCoord().toBigInteger()).mod(n);

		return expectedR.equals(r);
	}

	private byte[] digestDoFinal() {
		byte[] result = new byte[digest.getDigestSize()];
		digest.doFinal(result, 0);

		reset();

		return result;
	}

	private byte[] getZ(byte[] userID) {
		digest.reset();

		addUserID(digest, userID);

		addFieldElement(digest, ecParams.getCurve().getA());
		addFieldElement(digest, ecParams.getCurve().getB());
		addFieldElement(digest, ecParams.getG().getAffineXCoord());
		addFieldElement(digest, ecParams.getG().getAffineYCoord());
		addFieldElement(digest, pubPoint.getAffineXCoord());
		addFieldElement(digest, pubPoint.getAffineYCoord());

		byte[] result = new byte[digest.getDigestSize()];

		digest.doFinal(result, 0);

		return result;
	}

	private void addUserID(Digest digest, byte[] userID) {
		int len = userID.length * 8;
		digest.update((byte) (len >> 8 & 0xFF));
		digest.update((byte) (len & 0xFF));
		digest.update(userID, 0, userID.length);
	}

	private void addFieldElement(Digest digest, ECFieldElement v) {
		byte[] p = v.getEncoded();
		digest.update(p, 0, p.length);
	}

	protected ECMultiplier createBasePointMultiplier() {
		return new FixedPointCombMultiplier();
	}

	protected BigInteger calculateE(byte[] message) {
		return new BigInteger(1, message);
	}

	protected BigInteger[] derDecode(byte[] encoding) throws IOException {
		ASN1Sequence seq = ASN1Sequence.getInstance(ASN1Primitive.fromByteArray(encoding));
		if (seq.size() != 2) {
			return null;
		}

		BigInteger r = ASN1Integer.getInstance(seq.getObjectAt(0)).getValue();
		BigInteger s = ASN1Integer.getInstance(seq.getObjectAt(1)).getValue();

		byte[] expectedEncoding = derEncode(r, s);
		if (!Arrays.constantTimeAreEqual(expectedEncoding, encoding)) {
			return null;
		}

		return new BigInteger[] { r, s };
	}

	protected byte[] derEncode(BigInteger r, BigInteger s) throws IOException {

		ASN1EncodableVector v = new ASN1EncodableVector();
		v.add(new ASN1Integer(r));
		v.add(new ASN1Integer(s));
		return new DERSequence(v).getEncoded(ASN1Encoding.DER);
	}
}
